#!/usr/bin/env python

import os, sys
import time
from astropy.io import fits
import numpy as np
from scipy.stats import sigmaclip
from astropy.table import Table

def write_log(log_msg):

    curTime = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
    msg = f'[{curTime}] -- run getNoise.py and saved {log_msg}\n'
    if os.path.exists('noise_log.txt'):
        fp = open('noise_log.txt', 'a')
    else:
        fp = open('noise_log.txt', 'w')
    
    fp.write(msg)
    fp.close()

def remove_overscan(imgdata):
    print('> doing overscan correction ...')
    img = np.zeros((9232,9216), dtype=float)

    for i in range(4):
        img[0:4616, i*1152:(i+1)*1152] = imgdata[0:4616,(i*1250+27):((i+1)*1250-71)]
        img[4616:, i*1152:(i+1)*1152] = imgdata[4784:,(i*1250+27):((i+1)*1250-71)] 
        oc1 = np.mean(imgdata[0:4616,((i+1)*1250-71):((i+1)*1250)], axis=1).astype(float) 
        oc2 = np.mean(imgdata[4784:,((i+1)*1250-71):((i+1)*1250)], axis=1).astype(float)
#        print('oc1.shape: {}'.format(oc1.shape))
        for j in range(4616):
            img[j,i*1152:(i+1)*1152] = img[j,i*1152:(i+1)*1152] - oc1[j]
            img[4616+j,i*1152:(i+1)*1152] = img[4616+j,i*1152:(i+1)*1152]- oc2[j]
    for i in range(4,8):
        img[0:4616, i*1152:(i+1)*1152] = imgdata[0:4616,(i*1250+71):((i+1)*1250-27)]
        img[4616:, i*1152:(i+1)*1152] = imgdata[4784:,(i*1250+71):((i+1)*1250-27)]
        oc1 = np.mean(imgdata[0:4616,(i*1250):(i*1250+71)], axis=1).astype(float)
        oc2 = np.mean(imgdata[4784:,(i*1250):(i*1250+71)], axis=1).astype(float)
        for j in range(4616):
            img[j, i*1152:(i+1)*1152] = img[j, i*1152:(i+1)*1152] - oc1[j]
            img[4616+j,i*1152:(i+1)*1152] = img[4616+j,i*1152:(i+1)*1152] - oc2[j]

    print('> overscan correction done!')
    return img

def main():
    fname = sys.argv[1]
    imgdata = fits.getdata(fname).astype(float)
    img = remove_overscan(imgdata)
    
    os = []
    noise_clip_4_4 = []
    noise_clip_3_3 = []
    for ch in range(16):
#        print('> estimating rms for ch-{:02d}'.format(ch+1))
        os.append(ch+1)
        if ch < 8:
            ch_data = img[0:4616,ch*1152:(ch+1)*1152]
        else:
            ch_data = img[4616:, (15-ch)*1152:(16-ch)*1152]

        rms44 = np.std( sigmaclip(ch_data, 4, 4)[0] )
        rms33 = np.std( sigmaclip(ch_data, 3, 3)[0] )
        noise_clip_4_4.append( np.round(rms44,4) )
        noise_clip_3_3.append( np.round(rms33,4) )
    
    tab = Table()
    tab['OS'] = os
    tab['Noise_Clip_4_4'] = noise_clip_4_4
    tab['Noise_Clip_3_3'] = noise_clip_3_3
    print(tab)
    curTime = time.strftime("%Y-%m-%dT%H:%M:%S", time.localtime())

    tab_name = fname.replace('.fits','') + '_{}.tab'.format(curTime)
    tab.write(tab_name, format='ipac', overwrite=True)
    write_log(tab_name)

#    print('\n# Channel  noise(clip,4,4)  noise(clip,3,3)\n')
#    for i in range(16):
#        print('> OS-{:02d}   {:6.3f}   {:6.3f}'.format(i, noise_clip_4_4[i], noise_clip_3_3[i]))

if __name__ == '__main__':
    main()
